/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.up.kafka.productor;

import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.support.SendResult;
import org.springframework.lang.Nullable;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.util.concurrent.ListenableFutureCallback;

import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * @author: baiMl
 * @title: AbstractMessageProducer
 * @projectName: 海联技术中台V2.0
 * @depart 软件事业部
 * @description:
 * @date: 2022/4/24 14:23
 */
public abstract class AbstractMessageProducer<V> implements KafkaProducer<V> {
    protected KafkaTemplate<String, V> kafkaTemplate;
    private final Logger logger = LoggerFactory.getLogger(AbstractMessageProducer.class);

    public AbstractMessageProducer(KafkaTemplate<String, V> kafkaTemplate) {
        this.kafkaTemplate = kafkaTemplate;
    }

    /**
     * 获取Header，子类可根据实际情况添加需要的header
     * @return 返回Headers
     */
    protected abstract Headers getHeaders();

    @Override
    public void sendOnce(MessageRecord<V> record) {
        send(record);
    }

    @Override
    public SendResult<String, V> sendSync(MessageRecord<V> record) throws ExecutionException, InterruptedException {
        ListenableFuture<SendResult<String, V>> future = send(record);
        return future.get();
    }

    @Override
    public SendResult<String, V> sendSync(MessageRecord<V> record, long time, TimeUnit timeUnit) throws InterruptedException, ExecutionException, TimeoutException {
        ListenableFuture<SendResult<String, V>> future = send(record);
        return future.get(time, timeUnit);
    }


    protected ListenableFuture<SendResult<String, V>> send(MessageRecord<V> record) {
        if (logger.isInfoEnabled()) {
            logger.info("AbstractMessageProducer.send  record={}", record);
        }
        String key = record.getKey();
        Integer partition = record.getPartition();
        V message = record.getMessage();
        String topic = record.getTopic();
        Headers headers = getHeaders();
        if (headers != null) {
            ProducerRecord<String, V> producerRecord;
            if (partition != null) {
                producerRecord = new ProducerRecord<>(topic, partition, key, message);
            } else {
                producerRecord = new ProducerRecord<>(topic, key, message);
            }
            for (Header header : headers) {
                producerRecord.headers().add(header);
            }
            return kafkaTemplate.send(producerRecord);
        } else {
            if (partition != null) {
                return kafkaTemplate.send(topic, partition, key, message);
            } else {
                return kafkaTemplate.send(topic, key, message);
            }
        }
    }

    @Override
    public void sendAndCallBack(MessageRecord<V> record, MsgCallBack<V> callBack) {
        ListenableFuture<SendResult<String, V>> send = send(record);
        send.addCallback(new ListenableFutureCallback<SendResult<String, V>>() {
            @Override
            public void onFailure( Throwable throwable) {
                if (logger.isErrorEnabled()) {
                    logger.error("消息发送出现异常，异常信息:{}", throwable.getMessage());
                }
                callBack.doError(throwable);
            }

            @Override
            public void onSuccess(@Nullable SendResult<String, V> result) {
                if (logger.isInfoEnabled()) {
                    logger.info("消息发送成功，返回发送结果 result={}", result);
                }
                callBack.doSuccess(result);
            }
        });
    }

}
